package top.yudoge.dogeeditor.app.web

import org.springframework.beans.factory.annotation.Autowired
import org.springframework.http.MediaType
import org.springframework.web.bind.annotation.PostMapping
import org.springframework.web.bind.annotation.RequestBody
import org.springframework.web.bind.annotation.RequestMapping
import org.springframework.web.bind.annotation.RestController
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter
import top.yudoge.dogeeditor.app.query.ChatRewriteQuery
import top.yudoge.dogeeditor.app.query.SaveKnowledgeQuery
import top.yudoge.dogeeditor.app.utils.ChatHistoryUtils
import top.yudoge.dogeeditor.app.utils.TokenUtils
import top.yudoge.dogeeditor.commons.utils.convert
import top.yudoge.dogeeditor.commons.utils.toJsonStr
import top.yudoge.dogeeditor.commons.web.Response
import top.yudoge.dogeeditor.llm.memory.CacheSessionMemory
import top.yudoge.dogeeditor.llm.memory.SessionMemory
import top.yudoge.dogeeditor.llm.rewrite.RewriteCallback
import top.yudoge.dogeeditor.llm.rewrite.RewriteEvent
import top.yudoge.dogeeditor.llm.rewrite.RewriteEventType
import top.yudoge.dogeeditor.llm.rewrite.SQLGenerate2ChainNode
import top.yudoge.dogeeditor.llm.rewrite.SQLRewriter
import top.yudoge.dogeeditor.service.DeKnowledgeService

@RestController
@RequestMapping("/chat/new")
class NewChatController {


    @Autowired
    private lateinit var sqlRewriter: SQLRewriter

    @Autowired
    private lateinit var knowledgeService: DeKnowledgeService

    @Autowired
    private lateinit var sessionMemory: SessionMemory

    @PostMapping
    fun testNewSQLRewriter(
        @RequestBody query: ChatRewriteQuery,
    ) {
        sqlRewriter.rewrite(query.convert(), object : RewriteCallback {
            override fun onEvent(event: RewriteEvent) {
                println("event: ${event.type} data: ${event.data}")
            }
        })

    }

    @PostMapping(produces = [MediaType.TEXT_EVENT_STREAM_VALUE])
    fun chatForRewrite(
        @RequestBody query: ChatRewriteQuery,
    ): SseEmitter {
        val emitter = SseEmitter(1000 * 60 * 60)

        val base64TokenBuffer = TokenUtils.base64TokenBuffer()

        sqlRewriter.rewrite(query.convert(), object : RewriteCallback {
            override fun onEvent(event: RewriteEvent) {

                if (event.type == RewriteEventType.Reasoning) {
                    emitter.send(SseEmitter.event().name("reasoning").data(event.data as String))
                }

                if (event.type == RewriteEventType.Token) {
                    base64TokenBuffer.append(event.data as String) ?.let { tokenCanSendSafely ->
                        emitter.send(SseEmitter.event().name("token").data(tokenCanSendSafely))
                    }
                }

                if (event.type == RewriteEventType.Error) {
                    val error = event.data as Throwable
                    base64TokenBuffer.remaining() ?.let { remainingToken ->
                        emitter.send(SseEmitter.event().name("token").data(remainingToken))
                    }
                    emitter.send(SseEmitter.event().name("error").data(error.message ?: ""))
                    emitter.completeWithError(error)
                }

                if (event.type == RewriteEventType.Complete) {
                    base64TokenBuffer.remaining() ?.let { remainingToken ->
                        emitter.send(SseEmitter.event().name("token").data(remainingToken))
                    }
                    emitter.send(SseEmitter.event().name("complete").data(""))
                    emitter.complete()
                }


                println("event: ${event.type} data: ${event.data}")
            }
        })

        return emitter
    }

    @PostMapping("/save-knowledge")
    fun saveKnowledge(
        @RequestBody query: SaveKnowledgeQuery,
    ): Response<Boolean> {

        var sessionId = query.sessionId
        var listChatHistory = ChatHistoryUtils.listChatHistory(sessionId, sessionMemory)

        if (listChatHistory.isEmpty()) {
            return Response.success(false)
        }

        var jsonStr = listChatHistory.toJsonStr()
        knowledgeService.save(jsonStr)

        return Response.success(true)
    }

}